package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"syscall"

	"github.com/kr/pty"
	"golang.org/x/crypto/ssh/terminal"
)

func mountProc(ctx *context) {
	dir := path.Join(ctx.workDir, "proc")
	os.RemoveAll(dir)
	assert(os.MkdirAll(dir, 0555))
	assert(syscall.Mount("proc", dir, "proc", syscall.MS_NOEXEC|syscall.MS_NOSUID|syscall.MS_NODEV, ""))
	ctx.procMounted = true
}

func mountDev(ctx *context) {
	dir := path.Join(ctx.workDir, "dev")
	os.RemoveAll(dir)
	assert(os.MkdirAll(dir, 0755))
	assert(syscall.Mount("tmpfs", dir, "tmpfs", syscall.MS_NOSUID|syscall.MS_STRICTATIME, "mode=755"))
	ctx.devMounted = true
}

func mountSys(ctx *context) {
	dir := path.Join(ctx.workDir, "sys")
	os.RemoveAll(dir)
	assert(os.MkdirAll(dir, 0555))
	assert(syscall.Mount("sysfs", dir, "sysfs", syscall.MS_NOSUID|syscall.MS_NODEV|syscall.MS_NOEXEC|syscall.MS_RELATIME, ""))
	ctx.sysMounted = true
}

func mountPts(ctx *context) {
	dir := path.Join(ctx.workDir, "/dev/pts")
	os.RemoveAll(dir)
	assert(os.MkdirAll(dir, 0620))
	assert(syscall.Mount("devpts", dir, "devpts", syscall.MS_NOSUID|syscall.MS_NOEXEC|syscall.MS_RELATIME, "mode=620,ptmxmode=666,gid=5"))
	ctx.ptsMounted = true
}

func pivotRoot(newroot string) {
	putold := filepath.Join(newroot, "/.pivot_root")

	assert(syscall.Mount(newroot, newroot, "", syscall.MS_BIND|syscall.MS_REC, ""))

	assert(os.MkdirAll(putold, 0700))
	defer os.RemoveAll(putold)

	assert(syscall.PivotRoot(newroot, putold))

	assert(os.Chdir("/"))

	putold = "/.pivot_root"
	assert(syscall.Unmount(putold, syscall.MNT_DETACH))
}

func doRun() {
	fs := flag.NewFlagSet("_run", flag.ExitOnError)
	fs.BoolVar(&DEBUG, "debug", false, "show debug info")
	configDir := fs.String("c", "", "config dir")
	assert(fs.Parse(os.Args[2:]))

	var ctx context
	defer func() {
		if ctx.procMounted {
			syscall.Unmount("proc", 0)
		}
		if ctx.devMounted {
			syscall.Unmount("tmpfs", 0)
		}
		if ctx.sysMounted {
			syscall.Unmount("sysfs", 0)
		}
		if ctx.ptsMounted {
			syscall.Unmount("devpts", 0)
		}
	}()

	// load config
	f, err := os.Open(*configDir)
	assert(err)
	var config Config
	assert(json.NewDecoder(f).Decode(&config))
	ctx.workDir = config.ImageDir

	// change host name
	assert(syscall.Sethostname([]byte(config.HostName)))

	// show info
	if DEBUG {
		fmt.Println("Welcome to ltc")
		fmt.Printf("PID: %d\n", os.Getpid())
	}

	mountProc(&ctx)
	mountSys(&ctx)
	mountDev(&ctx)
	mountPts(&ctx)

	pivotRoot(ctx.workDir)

	os.Remove("/dev/null")
	os.Remove("/dev/zero")
	os.Remove("/dev/urandom")
	os.Remove("/dev/random")
	assert(syscall.Mknod("/dev/null", syscall.S_IFBLK, 259))
	assert(syscall.Chmod("/dev/null", 0666))
	assert(syscall.Mknod("/dev/zero", syscall.S_IFBLK, 261))
	assert(syscall.Chmod("/dev/zero", 0666))
	assert(syscall.Mknod("/dev/urandom", syscall.S_IFCHR, 265))
	assert(syscall.Chmod("/dev/urandom", 0666))
	assert(syscall.Mknod("/dev/random", syscall.S_IFCHR, 264))
	assert(syscall.Chmod("/dev/random", 0666))

	assert(os.Symlink("/dev/pts/ptmx", "/dev/ptmx"))
	assert(os.Chmod("/dev/ptmx", 0666))

	assert(syscall.Mknod("/dev/tty", syscall.S_IFCHR, 1280))

	mountVolumns(config.Mounts.Mounts, config.Mounts.Volumns)

	cmd := exec.Command(fs.Arg(0), fs.Args()[1:]...)
	cmd.Env = append(os.Environ(), defaultEnv...)
	cmd.Env = append(cmd.Env, "HOSTNAME="+config.HostName)
	cmd.Env = append(cmd.Env, "HOME=/root")

	if config.Tty {
		ptmx, err := pty.Start(cmd)
		assert(err)
		defer ptmx.Close()

		ch := make(chan os.Signal, 1)
		signal.Notify(ch, syscall.SIGWINCH)
		go func() {
			for range ch {
				assert(pty.InheritSize(os.Stdin, ptmx))
			}
		}()
		ch <- syscall.SIGWINCH

		old, err := terminal.MakeRaw(int(os.Stdin.Fd()))
		assert(err)
		defer terminal.Restore(int(os.Stdin.Fd()), old)

		go io.Copy(ptmx, os.Stdin)
		io.Copy(os.Stdout, ptmx)
	} else {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin
		assert(cmd.Run())
	}
}
